package cn.springcloud.alibaba.auth.config;

import cn.springcloud.alibaba.auth.common.TenantContextHolder;
import cn.springcloud.alibaba.core.base.Result;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.security.authentication.InsufficientAuthenticationException;
import org.springframework.security.core.Authentication;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.common.util.OAuth2Utils;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.stereotype.Component;

import java.security.Principal;
import java.util.Map;

/**
 * oauth-token拦截器，统一返回token格式
 */
@Slf4j
@Component
@Aspect
public class OauthTokenAspect {

	@Around("execution(* org.springframework.security.oauth2.provider.endpoint.TokenEndpoint.postAccessToken(..))")
	public Object handleControllerMethod(ProceedingJoinPoint joinPoint) throws Throwable {
		try {
			Object[] args = joinPoint.getArgs();
			Principal principal = (Principal) args[0];
			if (!(principal instanceof Authentication)) {
				throw new InsufficientAuthenticationException("There is no client authentication. Try adding an appropriate authentication filter.");
			}
			String clientId = this.getClientId(principal);
			Map<String, String> parameters = (Map<String, String>) args[1];
			String grantType = parameters.get(OAuth2Utils.GRANT_TYPE);

			TenantContextHolder.setTenantId(clientId);

			Object proceed = joinPoint.proceed();

			ResponseEntity<OAuth2AccessToken> responseEntity = (ResponseEntity<OAuth2AccessToken>) proceed;
			OAuth2AccessToken body = responseEntity.getBody();
			return ResponseEntity.status(HttpStatus.OK).body(Result.success(body));

			// 如果使用 @EnableOAuth2Sso 注解不能修改返回格式，否则授权码模式可以统一改
			/*if ("authorization_code".equals(grantType)) {
				return proceed;
			} else {
				ResponseEntity<OAuth2AccessToken> responseEntity = (ResponseEntity<OAuth2AccessToken>) proceed;
				OAuth2AccessToken body = responseEntity.getBody();
				return ResponseEntity.status(HttpStatus.OK).body(Result.success(body));
			}*/
		} finally {
			TenantContextHolder.clear();
		}
	}

	private String getClientId(Principal principal) {
		Authentication client = (Authentication) principal;
		if (!client.isAuthenticated()) {
			throw new InsufficientAuthenticationException("The client is not authenticated.");
		}
		String clientId = client.getName();
		if (client instanceof OAuth2Authentication) {
			clientId = ((OAuth2Authentication) client).getOAuth2Request().getClientId();
		}
		return clientId;
	}
}
